1794E - Labeling the Tree with Distances - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define REP(i,a,b) for(int i=(a);i<(b);i++)
#define RREP(i,a,b) for(int i=(a)-1;i>(b)-1;i--)
#define F first
#define S second
#define pb push_back
#define mp make_pair
#define zero(a) memset(a,0,sizeof(a))
#define neg(a) memset(a,-1,sizeof(a))
typedef vector<int> vi;
typedef vector<vector<int> > vvi;
typedef pair<int,int> pi;
typedef vector<pi> vpi;
typedef vector<long long> vll;
typedef long long ll;
typedef pair<double,double> pd;
typedef pair<long long, long long> pll;


int MOD=2567387659;
int norm(int x) {
    if (x < 0) {
        x += MOD;
    }
    if (x >= MOD) {
        x -= MOD;
    }
    return x;
}
 
template<class T>
T power(T a, ll b) {
    T res = 1;
    for (; b; b /= 2, a *= a) {
        if (b % 2) {
            res *= a;
        }
    }
    return res;
}
template<class T>
T inv(T a) {return power(a, MOD-2);}
struct Z {
    int x;
    Z(int x = 0) : x(norm(x)) {}
    int val() const {
        return x;
    }
    Z operator-() const {
        return Z(norm(MOD - x));
    }
    Z inv() const {
        assert(x != 0);
        return power(*this, MOD - 2);
    }
    Z &operator*=(const Z &rhs) {
        x = ll(x) * rhs.x % MOD;
        return *this;
    }
    Z &operator+=(const Z &rhs) {
        x = norm(x + rhs.x);
        return *this;
    }
    Z &operator-=(const Z &rhs) {
        x = norm(x - rhs.x);
        return *this;
    }
    Z &operator/=(const Z &rhs) {
        return *this *= rhs.inv();
    }
    friend Z operator*(const Z &lhs, const Z &rhs) {
        Z res = lhs;
        res *= rhs;
        return res;
    }
    friend Z operator+(const Z &lhs, const Z &rhs) {
        Z res = lhs;
        res += rhs;
        return res;
    }
    friend Z operator-(const Z &lhs, const Z &rhs) {
        Z res = lhs;
        res -= rhs;
        return res;
    }
    friend Z operator/(const Z &lhs, const Z &rhs) {
        Z res = lhs;
        res /= rhs;
        return res;
    }
};
Z B=1235447;
vector<vi> edge;
Z dp[2][200000];
int anc[200000];
Z dfs(int n){
    Z ans=0;
    for(auto x : edge[n]){
        if(x!=anc[n]){
            anc[x]=n;
            ans+=dfs(x);
        }
    }
    return dp[0][n]=ans*B+1;
}
void dfs2(int n){
    if(n!=0){
        dp[1][n]=(dp[1][anc[n]]+dp[0][anc[n]]-dp[0][n]*B)*B;
    }
    for(auto x : edge[n]){
        if(x!=anc[n]){
            dfs2(x);
        }
    }
}
Z c[200000];
signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    //freopen("in.txt", "r", stdin);
    //freopen("out.txt", "w", stdout);
    int n; cin >> n;
    int ar[n-1];REP(i,0,n-1)cin >> ar[i];
    edge.assign(n,{});
    int a,b;
    REP(i,0,n-1){
        cin >> a >> b;
        a--;b--;
        edge[a].pb(b);
        edge[b].pb(a);
    }
    dfs(0);
    dfs2(0);
    REP(i,0,n-1){
        c[ar[i]]+=1;
    }
    Z sum=0;
    REP(i,0,n){
        sum+=c[i]*power(B,i);
    }
    set<int> pos;
    Z d;
    REP(i,0,n){
        d=sum+power(B,i);
        pos.insert(d.val());
    }
    set<int> ans1;
    REP(i,0,n){
        d=dp[0][i]+dp[1][i];
        if(pos.count(d.val()))ans1.insert(i);
    }
    
    zero(c);
    pos.clear();
    MOD=1562386201;
    B=565867;
    
    dfs(0);
    dfs2(0);
    REP(i,0,n-1){
        c[ar[i]]+=1;
    }
    sum=0;
    REP(i,0,n){
        sum+=c[i]*power(B,i);
    }
    REP(i,0,n){
        d=sum+power(B,i);
        pos.insert(d.val());
    }
    set<int> ans2;
    REP(i,0,n){
        d=dp[0][i]+dp[1][i];
        if(pos.count(d.val()))ans2.insert(i);
    }
    set<int> tans;
    for(auto x : ans1){
        if(ans2.count(x))tans.insert(x);
    }
    cout << tans.size() << "\n";
    for(auto x : tans)cout << x+1 <<" ";cout << endl;
}